"""记忆管理器 - 多层记忆架构管理"""

import asyncio
from typing import Dict, Any, List, Optional, Union
from datetime import datetime, timedelta
from loguru import logger

from ..database.manager import DatabaseManager


class MemoryManager:
    """记忆管理器 - 管理短期、长期、关联和语义记忆"""
    
    def __init__(self, db_manager: DatabaseManager, config: Dict[str, Any]):
        self.db_manager = db_manager
        self.config = config
        self.memory_config = config.get('memory', {})
        
        # 记忆配置
        self.short_term_ttl = self.memory_config.get('short_term_ttl', 3600)  # 1小时
        self.long_term_threshold = self.memory_config.get('long_term_threshold', 5)  # 访问5次后转为长期记忆
        self.semantic_similarity_threshold = self.memory_config.get('semantic_similarity_threshold', 0.7)
        self.max_memories_per_user = self.memory_config.get('max_memories_per_user', 10000)
        
        # 初始化状态
        self.is_initialized = False
        
        logger.info("记忆管理器已创建")
    
    async def initialize(self):
        """初始化记忆管理器"""
        if self.is_initialized:
            logger.warning("记忆管理器已初始化")
            return
        
        try:
            # 确保数据库连接已初始化
            if not self.db_manager.is_initialized:
                await self.db_manager.initialize()
            
            self.is_initialized = True
            logger.info("记忆管理器初始化完成")
            
        except Exception as e:
            logger.error(f"初始化记忆管理器失败: {e}")
            raise
    
    async def shutdown(self):
        """关闭记忆管理器"""
        if not self.is_initialized:
            logger.warning("记忆管理器未初始化")
            return
        
        try:
            self.is_initialized = False
            logger.info("记忆管理器已关闭")
            
        except Exception as e:
            logger.error(f"关闭记忆管理器失败: {e}")
    
    async def create_user_context(self, user_id: str) -> Dict[str, Any]:
        """为用户创建记忆上下文"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 获取用户状态
            user_state = await self.db_manager.redis_client.get_user_state(user_id)
            
            # 获取最近的短期记忆
            recent_memories = await self.get_recent_memories(user_id, limit=10)
            
            # 获取用户统计信息
            user_stats = await self.db_manager.mongodb_client.get_user_stats(user_id)
            
            context = {
                'user_id': user_id,
                'user_state': user_state,
                'recent_memories': recent_memories,
                'user_stats': user_stats,
                'created_at': datetime.utcnow().isoformat()
            }
            
            logger.debug(f"创建用户记忆上下文: {user_id}")
            return context
            
        except Exception as e:
            logger.error(f"创建用户记忆上下文失败: {e}")
            raise
    
    async def store_memory(self, user_id: str, content: str, memory_type: str = 'general',
                          metadata: Optional[Dict[str, Any]] = None,
                          importance: float = 0.5) -> str:
        """存储记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            import uuid
            memory_id = str(uuid.uuid4())
            timestamp = datetime.utcnow()
            
            memory_data = {
                'memory_id': memory_id,
                'user_id': user_id,
                'content': content,
                'memory_type': memory_type,
                'importance': importance,
                'access_count': 0,
                'created_at': timestamp,
                'last_accessed': timestamp,
                'metadata': metadata or {}
            }
            
            # 1. 存储到短期记忆（Redis）
            await self.db_manager.redis_client.store_memory(
                user_id, memory_id, memory_data, ttl=self.short_term_ttl
            )
            
            # 2. 存储到长期记忆（MongoDB）
            await self.db_manager.mongodb_client.store_memory(memory_data)
            
            # 3. 存储到语义记忆（ChromaDB）
            await self.db_manager.chromadb_client.store_memory(
                user_id, content, memory_type, None, importance, memory_id
            )
            
            # 4. 创建关联记忆（Neo4j）
            await self._create_memory_associations(user_id, memory_id, content, memory_type)
            
            # 5. 更新用户统计
            await self._update_user_memory_stats(user_id)
            
            logger.debug(f"存储记忆: {memory_id} for user {user_id}")
            return memory_id
            
        except Exception as e:
            logger.error(f"存储记忆失败: {e}")
            raise
    
    async def _create_memory_associations(self, user_id: str, memory_id: str, 
                                        content: str, memory_type: str):
        """创建记忆关联"""
        try:
            # 确保用户节点存在
            await self.db_manager.neo4j_client.create_user_node(user_id, {'user_id': user_id})
            
            # 创建记忆节点
            await self.db_manager.neo4j_client.create_memory_node(
                memory_id, {
                    'content': content,
                    'memory_type': memory_type,
                    'user_id': user_id
                }
            )
            
            # 提取关键概念并创建关联
            concepts = await self._extract_concepts(content)
            for concept in concepts:
                concept_id = f"concept_{hash(concept) % 1000000}"
                await self.db_manager.neo4j_client.create_concept_node(
                    concept_id, {'name': concept, 'type': 'extracted'}
                )
                await self.db_manager.neo4j_client.create_relationship(
                    memory_id, 'Memory', concept_id, 'Concept', 'RELATES_TO', {'strength': 1.0}
                )
            
            # 查找相似记忆并创建关联
            similar_memories = await self.search_memories(
                user_id, content, limit=3, exclude_memory_id=memory_id
            )
            
            for similar_memory in similar_memories:
                if similar_memory.get('similarity', 0) > self.semantic_similarity_threshold:
                    await self.db_manager.neo4j_client.create_relationship(
                        memory_id, 'Memory', similar_memory['memory_id'], 'Memory', 'SIMILAR_TO',
                        {'similarity': similar_memory['similarity']}
                    )
            
        except Exception as e:
            logger.error(f"创建记忆关联失败: {e}")
    
    async def _extract_concepts(self, content: str) -> List[str]:
        """从内容中提取关键概念（简化版本）"""
        try:
            # 这里是一个简化的概念提取逻辑
            # 在实际应用中，可以使用NLP库或LLM来提取概念
            
            import re
            
            # 简单的关键词提取
            words = re.findall(r'\b\w{3,}\b', content.lower())
            
            # 过滤常见停用词
            stop_words = {'的', '了', '在', '是', '我', '你', '他', '她', '它', '我们', '你们', '他们',
                         'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for',
                         'of', 'with', 'by', 'from', 'up', 'about', 'into', 'through', 'during'}
            
            concepts = [word for word in set(words) if word not in stop_words and len(word) > 2]
            
            # 限制概念数量
            return concepts[:10]
            
        except Exception as e:
            logger.error(f"提取概念失败: {e}")
            return []
    
    async def get_memory(self, memory_id: str, user_id: str) -> Optional[Dict[str, Any]]:
        """获取特定记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 首先从短期记忆（Redis）查找
            memory = await self.db_manager.redis_client.get_memory(user_id, memory_id)
            
            if not memory:
                # 从长期记忆（MongoDB）查找 - 使用正确的方法
                # MongoDB客户端没有get_memory方法，需要通过查询memories集合来获取
                collection = self.db_manager.mongodb_client.get_collection('memories')
                memory_doc = await collection.find_one({"memory_id": memory_id, "user_id": user_id})
                if memory_doc:
                    memory_doc['_id'] = str(memory_doc['_id'])
                    memory = memory_doc
            
            if memory:
                # 更新访问计数和时间
                await self._update_memory_access(memory_id, user_id)
            
            return memory
            
        except Exception as e:
            logger.error(f"获取记忆失败: {e}")
            return None
    
    async def _update_memory_access(self, memory_id: str, user_id: str):
        """更新记忆访问信息"""
        try:
            current_time = datetime.utcnow()
            
            # 更新MongoDB中的访问信息 - 使用正确的方法
            collection = self.db_manager.mongodb_client.get_collection('memories')
            await collection.update_one(
                {"memory_id": memory_id, "user_id": user_id},
                {
                    "$set": {"last_accessed": current_time},
                    "$inc": {"access_count": 1}
                }
            )
            
            # 获取访问计数 - 使用正确的方法
            memory_doc = await collection.find_one({"memory_id": memory_id, "user_id": user_id})
            if memory_doc and memory_doc.get('access_count', 0) >= self.long_term_threshold:
                # 如果访问次数达到阈值，延长Redis中的TTL
                await self.db_manager.redis_client.extend_memory_ttl(
                    user_id, memory_id, self.short_term_ttl * 2
                )
            
        except Exception as e:
            logger.error(f"更新记忆访问失败: {e}")
    
    async def search_memories(self, user_id: str, query: str, 
                            memory_types: Optional[List[str]] = None,
                            limit: int = 10, exclude_memory_id: Optional[str] = None) -> List[Dict[str, Any]]:
        """搜索用户记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 1. 语义搜索（ChromaDB）
            semantic_results = await self.db_manager.chromadb_client.search_memories(
                user_id, query, memory_type=None, n_results=limit * 2
            )
            
            # 2. 关联搜索（Neo4j）
            related_memories = await self.db_manager.neo4j_client.find_related_memories(
                query, user_id, limit=limit
            )
            
            # 3. 合并和排序结果
            all_results = {}
            
            # 添加语义搜索结果
            for result in semantic_results:
                # ChromaDB返回的是'id'字段，需要映射为'memory_id'
                memory_id = result.get('id') or result.get('memory_id')
                if memory_id and memory_id != exclude_memory_id:
                    # 确保结果中包含memory_id字段
                    result_with_memory_id = {
                        **result,
                        'memory_id': memory_id,
                        'semantic_score': result.get('similarity', 0),
                        'relation_score': 0
                    }
                    all_results[memory_id] = result_with_memory_id
            
            # 添加关联搜索结果
            for result in related_memories:
                memory_id = result.get('memory_id')
                if memory_id and memory_id != exclude_memory_id:
                    if memory_id in all_results:
                        all_results[memory_id]['relation_score'] = result.get('relevance', 0)
                    else:
                        all_results[memory_id] = {
                            **result,
                            'semantic_score': 0,
                            'relation_score': result.get('relevance', 0)
                        }
            
            # 4. 计算综合得分并排序
            for memory_id, result in all_results.items():
                semantic_score = result.get('semantic_score', 0)
                relation_score = result.get('relation_score', 0)
                
                # 综合得分：语义相似度 * 0.7 + 关联度 * 0.3
                combined_score = semantic_score * 0.7 + relation_score * 0.3
                result['combined_score'] = combined_score
            
            # 按综合得分排序
            sorted_results = sorted(
                all_results.values(),
                key=lambda x: x.get('combined_score', 0),
                reverse=True
            )
            
            # 5. 过滤记忆类型
            if memory_types:
                sorted_results = [
                    result for result in sorted_results
                    if result.get('memory_type') in memory_types
                ]
            
            # 6. 获取完整记忆信息
            final_results = []
            for result in sorted_results[:limit]:
                memory_id = result.get('memory_id')
                if not memory_id:
                    continue

                full_memory = await self.get_memory(memory_id, user_id)
                if full_memory:
                    # 如果在缓存/数据库中找到了完整记忆，合并评分信息后返回
                    full_memory.update({
                        'similarity': result.get('combined_score', 0),
                        'semantic_score': result.get('semantic_score', 0),
                        'relation_score': result.get('relation_score', 0)
                    })
                    final_results.append(full_memory)
                else:
                    # 无法获取完整记忆时，直接返回语义搜索结果作为降级处理
                    fallback_memory = {
                        'memory_id': memory_id,
                        'content': result.get('content') or result.get('doc') or '',
                        'metadata': result.get('metadata') or result.get('metadatas') or {},
                        'similarity': result.get('combined_score', 0),
                        'semantic_score': result.get('semantic_score', 0),
                        'relation_score': result.get('relation_score', 0)
                    }
                    final_results.append(fallback_memory)
            
            logger.debug(f"搜索记忆: {user_id} - 找到 {len(final_results)} 条结果")
            return final_results
            
        except Exception as e:
            logger.error(f"搜索记忆失败: {e}")
            return []
    
    async def get_recent_memories(self, user_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """获取用户最近的记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 从Redis获取最近的短期记忆
            recent_memories = await self.db_manager.redis_client.get_recent_memories(
                user_id, limit
            )
            
            # 如果短期记忆不足，从MongoDB补充
            if len(recent_memories) < limit:
                additional_needed = limit - len(recent_memories)
                long_term_memories = await self.db_manager.mongodb_client.get_recent_memories(
                    user_id, additional_needed
                )
                recent_memories.extend(long_term_memories)
            
            return recent_memories[:limit]
            
        except Exception as e:
            logger.error(f"获取最近记忆失败: {e}")
            return []
    
    async def delete_memory(self, memory_id: str, user_id: str) -> bool:
        """删除记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 从所有存储层删除
            tasks = [
                self.db_manager.redis_client.delete_memory(user_id, memory_id),
                self.db_manager.mongodb_client.delete_memory(memory_id),
                self.db_manager.chromadb_client.delete_memory(memory_id),
                self.db_manager.neo4j_client.delete_memory_node(memory_id)
            ]
            
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            # 检查是否有错误
            success = all(not isinstance(result, Exception) for result in results)
            
            if success:
                logger.info(f"删除记忆: {memory_id} for user {user_id}")
            else:
                logger.warning(f"删除记忆部分失败: {memory_id}")
            
            return success
            
        except Exception as e:
            logger.error(f"删除记忆失败: {e}")
            return False
    
    async def update_user_state(self, user_id: str, state_updates: Dict[str, Any]):
        """更新用户状态"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            await self.db_manager.redis_client.update_user_state(user_id, state_updates)
            logger.debug(f"更新用户状态: {user_id}")
            
        except Exception as e:
            logger.error(f"更新用户状态失败: {e}")
    
    async def get_user_context(self, user_id: str) -> Dict[str, Any]:
        """获取用户记忆上下文"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            # 获取用户状态
            user_state = await self.db_manager.redis_client.get_user_state(user_id)
            
            # 获取记忆统计
            memory_stats = await self.db_manager.mongodb_client.get_user_stats(user_id)
            
            # 获取最近记忆
            recent_memories = await self.get_recent_memories(user_id, limit=5)
            
            # 获取关联图统计
            graph_stats = await self.db_manager.neo4j_client.get_user_graph_stats(user_id)
            
            return {
                'user_id': user_id,
                'user_state': user_state,
                'memory_stats': memory_stats,
                'recent_memories': recent_memories,
                'graph_stats': graph_stats,
                'context_created_at': datetime.utcnow().isoformat()
            }
            
        except Exception as e:
            logger.error(f"获取用户记忆上下文失败: {e}")
            return {'user_id': user_id, 'error': str(e)}
    
    async def _update_user_memory_stats(self, user_id: str):
        """更新用户记忆统计"""
        try:
            current_time = datetime.utcnow()
            
            # 更新MongoDB中的用户统计
            await self.db_manager.mongodb_client.update_user_stats(
                user_id, {
                    'last_memory_created': current_time,
                    'total_memories': {'$inc': 1}
                }
            )
            
        except Exception as e:
            logger.error(f"更新用户记忆统计失败: {e}")
    
    async def cleanup_old_memories(self, days_threshold: int = 30) -> int:
        """清理旧记忆"""
        if not self.is_initialized:
            raise RuntimeError("记忆管理器未初始化")
        
        try:
            cutoff_date = datetime.utcnow() - timedelta(days=days_threshold)
            
            # 清理各个存储层的旧记忆
            tasks = [
                self.db_manager.mongodb_client.cleanup_old_memories(cutoff_date),
                self.db_manager.chromadb_client.cleanup_old_vectors(cutoff_date),
                self.db_manager.neo4j_client.cleanup_old_memories(cutoff_date)
            ]
            
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            total_cleaned = sum(result for result in results if isinstance(result, int))
            
            logger.info(f"清理旧记忆完成: 删除了 {total_cleaned} 条记忆")
            return total_cleaned
            
        except Exception as e:
            logger.error(f"清理旧记忆失败: {e}")
            return 0
    
    async def get_stats(self) -> Dict[str, Any]:
        """获取记忆管理器统计信息"""
        try:
            if not self.is_initialized:
                return {'status': 'not_initialized'}
            
            # 获取各存储层统计
            redis_stats = await self.db_manager.redis_client.get_stats()
            mongodb_stats = await self.db_manager.mongodb_client.get_stats()
            chromadb_stats = await self.db_manager.chromadb_client.get_stats()
            neo4j_stats = await self.db_manager.neo4j_client.get_stats()
            
            return {
                'status': 'initialized',
                'config': {
                    'short_term_ttl': self.short_term_ttl,
                    'long_term_threshold': self.long_term_threshold,
                    'semantic_similarity_threshold': self.semantic_similarity_threshold,
                    'max_memories_per_user': self.max_memories_per_user
                },
                'storage_stats': {
                    'redis': redis_stats,
                    'mongodb': mongodb_stats,
                    'chromadb': chromadb_stats,
                    'neo4j': neo4j_stats
                }
            }
            
        except Exception as e:
            logger.error(f"获取记忆管理器统计失败: {e}")
            return {'error': str(e)}
    
    async def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        try:
            # 检查初始化状态
            initialized = self.is_initialized
            
            # 检查数据库连接
            db_health = await self.db_manager.health_check()
            db_healthy = all(db_health.values())
            
            overall_healthy = initialized and db_healthy
            
            return {
                'healthy': overall_healthy,
                'status': 'healthy' if overall_healthy else 'unhealthy',
                'checks': {
                    'initialized': initialized,
                    'database': db_healthy
                },
                'database_health': db_health
            }
            
        except Exception as e:
            logger.error(f"记忆管理器健康检查失败: {e}")
            return {
                'healthy': False,
                'status': 'error',
                'error': str(e)
            }